# author@bupt_kt
# date@09/04/2021


import numpy as np
data=[]
# data[i][0]:第i个节点下的数据样例点，data[i][1]:第i个节点下的信息：该节点的预测值，该节点下一步的分类维度，左节点下标，右节点下标

def input_data():
    sample=[]
    file=open(r'Mnist/mnist_train.csv','r')
    i=0
    for line in file.readlines():
        l=line.strip().split(',')
        sample.append([])
        sample[i].append(int(l[0]))
        for num in l[1:]:
            if int(num) < 128:
                sample[i].append(0)
            else:
                sample[i].append(1)
        i+=1

    return sample

# 采用递归方式构造决策树
def decision(dimension,index):   ## dimension表示当前分支下能使用的维度分类，index表示当前节点下标
    ## 计算信息增益
    H = [0]*785  #信息增益
    H[0]=-1
    y = 10 * [0] #y[i]:数据中y=i的个数
    x = [2*[0] for i in range(785)]  #x[i][j]:数据中第i个x等于j的个数
    xy = []  #xy[i][j][k]:数据中第i个x等于j且y等于k的个数
    for i in range(785):
        xy.append([])
        for j in range(2):
            xy[i].append([0 for k in range(10)])

    for each in data[index][0]:
        y[each[0]] += 1
        for j in range(1,len(each)):
            x[j][each[j]] += 1
            xy[j][each[j]][each[0]] += 1

    hy=0
    for j in range(10):
        if y[j] != 0:
            hy += -1 * (y[j] / len(data[index][0])) * np.log2(y[j] / len(data[index][0]))
    for i in range(1,785):
        if dimension[i]==-1:  ##该维度已经使用过
            H[i]=-1
        else:
            hxy=0
            for j in range(2):
                for k in range(10):
                    if x[i][j]!=0 and xy[i][j][k]!=0:
                        hxy+=-1*(x[i][j]/len(data[index][0]))*(xy[i][j][k]/x[i][j])*np.log2(xy[i][j][k]/x[i][j])
            H[i]=hy-hxy

    info=[] # 添加节点信息
    info.append(y.index(max(y)))
    print(y.index(max(y)),y)
    if max(H)<0.1 or max(y)==len(data[index][0]): #信息增益阈值，退出条件
        info.append(-1)
        info.append(-1)
        info.append(-1)
        data[index].append(info)
        return
    else:
        info.append(H.index(max(H)))
        info.append(len(data))
        info.append(len(data)+1)
        dimension[info[1]]=-1
        left=info[2]
        right=info[3]
        data.append([])
        data.append([])
        data[left].append([])
        data[right].append([])
        for each in data[index][0]:
            if each[info[1]]==0:
                data[left][0].append(each)
            else:
                data[right][0].append(each)
        data[index].append(info)
        decision(dimension,left)
        decision(dimension,right)
        return


def test():#模型测试
    sample = []
    file = open(r'Mnist/mnist_test.csv', 'r')
    i = 0
    for line in file.readlines():
        l = line.strip().split(',')
        sample.append([])
        sample[i].append(int(l[0]))
        for num in l[1:]:
            if int(num) < 128:
                sample[i].append(0)
            else:
                sample[i].append(1)
        i += 1

    x=0
    error=0
    for each in sample:
        while data[x][1][2]!=-1:
            if each[data[x][1][1]] == 0:
                x=data[x][1][2]
            else:
                x=data[x][1][3]

        predict=data[x][1][0]
        if predict!=each[0]:
            error+=1

    print(1-error/len(sample))

if __name__ == '__main__':
    data.append([])
    dimension=785*[0]
    dimension[0]=-1
    index=0
    data[0].append(input_data())
    decision(dimension,index)
    test()